package com.suxin.ratelimit.repository;

import com.suxin.ratelimit.Rate;
import com.suxin.ratelimit.properties.RateLimiterProperties;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.ValueOperations;

import java.util.concurrent.TimeUnit;

/**
 * @author Tang
 * @classname RedisRateLimiter
 * @description [ redis 方式实现限流 ]
 * @date 2021/7/14 22:48
 */
public class RedisRateLimiter implements RateLimiter {

    private RedisTemplate<String, Object> redisTemplate;

    public RedisRateLimiter(RedisTemplate<String, Object> redisTemplate) {
        this.redisTemplate = redisTemplate;
    }

    @Override
    public Rate consumer(RateLimiterProperties.Policy policy, String key) {
        ValueOperations<String, Object> opsForValue = redisTemplate.opsForValue();
        Long limit = policy.getLimit();
        // 尝试设置 key
        Boolean absent = opsForValue.setIfAbsent(key, limit, policy.getRefreshInterval(), TimeUnit.SECONDS);
        // 设置失败
        if (!absent) {
            // 增加访问次数
            Integer count = Integer.valueOf(opsForValue.get(key).toString());
            if (count != null && count >= 0) {
                opsForValue.increment(key, -1);
            }
            return new Rate(null,key,policy.getRefreshInterval(), (long) (count - 1),System.currentTimeMillis());
        }
        return new Rate(null,key,policy.getRefreshInterval(),limit,System.currentTimeMillis());
    }
}